⚡️ Speed up function `_get_constraint_vals_and_feasibility` by 9% #38

codeflash-ai · 2025-10-29T02:57:55Z

📄 9% (0.09x) speedup for `_get_constraint_vals_and_feasibility` in `optuna/samplers/_gp/sampler.py`

⏱️ Runtime : 1.61 milliseconds → 1.48 milliseconds (best of 146 runs)

📝 Explanation and details

The optimization replaces the O(n²) list comprehension with validation approach with a single O(n) loop that builds the constraint list and validates lengths incrementally.

Key changes:

Single-pass processing: Instead of creating the full _constraint_vals list first and then validating all lengths with any(), the optimized version processes trials one by one, checking constraint length consistency as it builds the list.
Early validation: The length check happens immediately when processing each trial, allowing for early exit on the first inconsistent constraint length rather than processing all trials before validation.
Eliminated redundant iteration: The original code iterates through trials twice - once in the list comprehension and once in the any() validation. The optimized version does both operations in a single loop.

Performance benefits:

Normal cases (4-11% faster): Eliminates the overhead of the any() generator expression and reduces total iterations from 2n to n
Error cases with inconsistent constraints (60-2261% faster): Massive speedup because the function can exit immediately upon finding the first constraint length mismatch, rather than processing all trials first

The optimization is particularly effective for large-scale scenarios where constraint length validation fails early, as shown in the test cases where inconsistent constraint detection becomes dramatically faster.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 31 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import numpy as np
# imports
import pytest
from optuna.samplers._gp.sampler import _get_constraint_vals_and_feasibility

# Constants and mocks to simulate Optuna's environment
_CONSTRAINTS_KEY = "constraints"

class DummyFrozenTrial:
    """Minimal mock of optuna.trial.FrozenTrial for testing."""
    def __init__(self, trial_id):
        self._trial_id = trial_id

class DummyStorage:
    """Mock of optuna's storage system for testing."""
    def __init__(self, system_attrs_by_id):
        # system_attrs_by_id: dict[int, dict[str, Any]]
        self._system_attrs_by_id = system_attrs_by_id

    def get_trial_system_attrs(self, trial_id):
        # Simulate KeyError if trial_id is missing
        return self._system_attrs_by_id[trial_id]

class DummyStudy:
    """Mock of optuna.study.Study for testing."""
    def __init__(self, storage):
        self._storage = storage
from optuna.samplers._gp.sampler import _get_constraint_vals_and_feasibility

# -----------------------
# Unit tests start here
# -----------------------

# 1. Basic Test Cases

def test_all_trials_feasible():
    """All trials have constraints <= 0 (all feasible)."""
    system_attrs = {
        1: {_CONSTRAINTS_KEY: (0.0, -1.0)},
        2: {_CONSTRAINTS_KEY: (-2.0, 0.0)},
        3: {_CONSTRAINTS_KEY: (-1.0, -1.0)},
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(1), DummyFrozenTrial(2), DummyFrozenTrial(3)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 31.9μs -> 30.6μs (4.09% faster)

def test_some_trials_infeasible():
    """Some trials have constraints > 0 (infeasible)."""
    system_attrs = {
        1: {_CONSTRAINTS_KEY: (1.0, -1.0)},
        2: {_CONSTRAINTS_KEY: (0.0, 0.0)},
        3: {_CONSTRAINTS_KEY: (-1.0, 2.0)},
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(1), DummyFrozenTrial(2), DummyFrozenTrial(3)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 20.9μs -> 19.6μs (6.60% faster)

def test_all_trials_infeasible():
    """All trials have at least one constraint > 0 (all infeasible)."""
    system_attrs = {
        1: {_CONSTRAINTS_KEY: (1.0, 2.0)},
        2: {_CONSTRAINTS_KEY: (0.1, 0.1)},
        3: {_CONSTRAINTS_KEY: (100.0, -1.0)},
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(1), DummyFrozenTrial(2), DummyFrozenTrial(3)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 19.7μs -> 18.1μs (8.70% faster)


def test_empty_constraints_tuple():
    """Trials present but constraints tuple is empty for all."""
    system_attrs = {
        1: {_CONSTRAINTS_KEY: ()},
        2: {_CONSTRAINTS_KEY: ()},
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(1), DummyFrozenTrial(2)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 33.2μs -> 30.5μs (8.78% faster)

# 2. Edge Test Cases

def test_missing_constraints_key():
    """Constraints key missing in system_attrs: should default to empty tuple."""
    system_attrs = {
        1: {},  # No constraints key
        2: {},  # No constraints key
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(1), DummyFrozenTrial(2)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 21.1μs -> 19.4μs (8.99% faster)

def test_inconsistent_constraints_length_raises():
    """If trials have different number of constraints, should raise ValueError."""
    system_attrs = {
        1: {_CONSTRAINTS_KEY: (1.0, 2.0)},
        2: {_CONSTRAINTS_KEY: (3.0,)},  # Only one constraint
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(1), DummyFrozenTrial(2)]
    with pytest.raises(ValueError, match="The number of constraints must be the same for all trials."):
        _get_constraint_vals_and_feasibility(study, trials) # 3.50μs -> 2.18μs (60.3% faster)

def test_trial_id_not_in_storage_raises():
    """If a trial_id is missing in storage, should raise KeyError."""
    system_attrs = {
        1: {_CONSTRAINTS_KEY: (1.0, 2.0)},
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(1), DummyFrozenTrial(2)]  # 2 is missing
    with pytest.raises(KeyError):
        _get_constraint_vals_and_feasibility(study, trials) # 2.19μs -> 2.06μs (6.33% faster)

def test_zero_constraints_all_feasible():
    """Trials with zero constraints (empty tuple) are always feasible."""
    system_attrs = {
        1: {_CONSTRAINTS_KEY: ()},
        2: {_CONSTRAINTS_KEY: ()},
        3: {_CONSTRAINTS_KEY: ()},
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(1), DummyFrozenTrial(2), DummyFrozenTrial(3)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 25.5μs -> 23.4μs (8.93% faster)

def test_single_trial():
    """Single trial, with constraints."""
    system_attrs = {
        42: {_CONSTRAINTS_KEY: (0.0, 1.0, -1.0)},
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(42)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 21.4μs -> 19.4μs (9.92% faster)

def test_all_zeros_constraints():
    """All constraints are exactly zero: should be feasible."""
    system_attrs = {
        1: {_CONSTRAINTS_KEY: (0.0, 0.0)},
        2: {_CONSTRAINTS_KEY: (0.0, 0.0)},
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(1), DummyFrozenTrial(2)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 20.4μs -> 18.5μs (10.6% faster)

# 3. Large Scale Test Cases

def test_large_number_of_trials_and_constraints():
    """Test with 500 trials and 10 constraints each."""
    n_trials = 500
    n_constraints = 10
    system_attrs = {
        i: {_CONSTRAINTS_KEY: tuple(float(i % 3 - 1) for _ in range(n_constraints))}
        for i in range(n_trials)
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(i) for i in range(n_trials)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 279μs -> 271μs (2.64% faster)
    # For i % 3 == 0 -> -1, feasible; i % 3 == 1 -> 0, feasible; i % 3 == 2 -> 1, infeasible
    expected_feasible = np.array([(i % 3 != 2) for i in range(n_trials)])

def test_large_number_of_constraints():
    """Test with 5 trials and 999 constraints each."""
    n_trials = 5
    n_constraints = 999
    system_attrs = {
        i: {_CONSTRAINTS_KEY: tuple(0.0 if (i+j)%2==0 else -1.0 for j in range(n_constraints))}
        for i in range(n_trials)
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(i) for i in range(n_trials)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 138μs -> 135μs (1.65% faster)

def test_large_scale_infeasible():
    """Test with 100 trials, each with 20 constraints, all constraints positive (all infeasible)."""
    n_trials = 100
    n_constraints = 20
    system_attrs = {
        i: {_CONSTRAINTS_KEY: tuple(1.0 for _ in range(n_constraints))}
        for i in range(n_trials)
    }
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(i) for i in range(n_trials)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 96.2μs -> 94.1μs (2.27% faster)

def test_large_scale_mixed_feasibility():
    """Test with 100 trials, each with 5 constraints, some feasible, some not."""
    n_trials = 100
    n_constraints = 5
    system_attrs = {}
    for i in range(n_trials):
        # Odd trials infeasible (one constraint positive), even trials feasible (all <= 0)
        if i % 2 == 0:
            constraints = tuple(-1.0 for _ in range(n_constraints))
        else:
            constraints = tuple(1.0 if j == 0 else -1.0 for j in range(n_constraints))
        system_attrs[i] = {_CONSTRAINTS_KEY: constraints}
    storage = DummyStorage(system_attrs)
    study = DummyStudy(storage)
    trials = [DummyFrozenTrial(i) for i in range(n_trials)]
    constraint_vals, is_feasible = _get_constraint_vals_and_feasibility(study, trials) # 58.7μs -> 57.4μs (2.16% faster)
    expected_feasible = np.array([(i % 2 == 0) for i in range(n_trials)])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import pytest  # used for our unit tests
from optuna.samplers._gp.sampler import _get_constraint_vals_and_feasibility

# --- Function and minimal dependencies for testing ---

# Minimal FrozenTrial class
class FrozenTrial:
    def __init__(self, trial_id):
        self._trial_id = trial_id

# Minimal Study and Storage classes for testing
_CONSTRAINTS_KEY = "constraints"

class DummyStorage:
    def __init__(self, trial_system_attrs):
        # trial_system_attrs: dict of trial_id -> system_attrs dict
        self._trial_system_attrs = trial_system_attrs

    def get_trial_system_attrs(self, trial_id):
        # Return system attrs for the given trial_id
        if trial_id not in self._trial_system_attrs:
            raise KeyError(f"Trial id {trial_id} not found.")
        return self._trial_system_attrs[trial_id]

class DummyStudy:
    def __init__(self, trial_system_attrs):
        self._storage = DummyStorage(trial_system_attrs)
from optuna.samplers._gp.sampler import _get_constraint_vals_and_feasibility

# --- Unit tests ---

# Basic Test Cases

def test_single_trial_feasible():
    # One trial, all constraints <= 0
    attrs = {0: {_CONSTRAINTS_KEY: (0, -1, 0)}}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 24.9μs -> 22.1μs (13.1% faster)

def test_single_trial_infeasible():
    # One trial, some constraints > 0
    attrs = {0: {_CONSTRAINTS_KEY: (0, 1, -1)}}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 19.6μs -> 17.8μs (10.4% faster)

def test_multiple_trials_mixed_feasibility():
    # Multiple trials, some feasible, some not
    attrs = {
        0: {_CONSTRAINTS_KEY: (0, -1)},
        1: {_CONSTRAINTS_KEY: (1, -1)},
        2: {_CONSTRAINTS_KEY: (0, 0)},
    }
    study = DummyStudy(attrs)
    trials = [FrozenTrial(i) for i in range(3)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 19.7μs -> 18.5μs (6.42% faster)

def test_empty_constraints_tuple():
    # Constraints tuple is empty
    attrs = {0: {_CONSTRAINTS_KEY: ()}}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 17.4μs -> 15.8μs (10.2% faster)

def test_missing_constraints_key_defaults_to_empty():
    # Constraints key missing, should default to empty tuple
    attrs = {0: {}}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 16.6μs -> 14.2μs (16.9% faster)

def test_all_trials_feasible():
    # All trials feasible
    attrs = {
        0: {_CONSTRAINTS_KEY: (0, -1, -2)},
        1: {_CONSTRAINTS_KEY: (-5, -6, 0)},
    }
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0), FrozenTrial(1)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 19.0μs -> 17.6μs (8.22% faster)

def test_all_trials_infeasible():
    # All trials infeasible
    attrs = {
        0: {_CONSTRAINTS_KEY: (1, 2)},
        1: {_CONSTRAINTS_KEY: (3, 4)},
    }
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0), FrozenTrial(1)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 17.8μs -> 15.9μs (11.9% faster)

# Edge Test Cases

def test_inconsistent_num_constraints_raises():
    # Trials have different numbers of constraints
    attrs = {
        0: {_CONSTRAINTS_KEY: (0, -1)},
        1: {_CONSTRAINTS_KEY: (1,)},
    }
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0), FrozenTrial(1)]
    with pytest.raises(ValueError, match="The number of constraints must be the same for all trials."):
        _get_constraint_vals_and_feasibility(study, trials) # 3.52μs -> 2.14μs (64.6% faster)

def test_trial_not_found_raises():
    # Trial id not present in storage
    attrs = {0: {_CONSTRAINTS_KEY: (0, -1)}}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(1)]  # id 1 missing
    with pytest.raises(KeyError):
        _get_constraint_vals_and_feasibility(study, trials) # 2.17μs -> 1.86μs (16.8% faster)




def test_zero_constraint_value():
    # Constraint value exactly zero
    attrs = {0: {_CONSTRAINTS_KEY: (0,)}}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 34.1μs -> 31.2μs (9.25% faster)

def test_large_negative_constraint_value():
    # Large negative constraint value
    attrs = {0: {_CONSTRAINTS_KEY: (-99999999,)}}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 21.0μs -> 18.8μs (11.8% faster)

def test_large_positive_constraint_value():
    # Large positive constraint value
    attrs = {0: {_CONSTRAINTS_KEY: (99999999,)}}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 18.9μs -> 17.2μs (9.93% faster)

# Large Scale Test Cases

def test_many_trials_and_constraints_all_feasible():
    # 100 trials, each with 10 constraints, all <= 0
    num_trials = 100
    num_constraints = 10
    attrs = {
        i: {_CONSTRAINTS_KEY: tuple(-j for j in range(num_constraints))}
        for i in range(num_trials)
    }
    study = DummyStudy(attrs)
    trials = [FrozenTrial(i) for i in range(num_trials)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 78.2μs -> 76.5μs (2.12% faster)
    for v in vals:
        pass

def test_many_trials_and_constraints_some_infeasible():
    # 100 trials, each with 10 constraints, every 5th trial is infeasible
    num_trials = 100
    num_constraints = 10
    attrs = {}
    for i in range(num_trials):
        if i % 5 == 0:
            # Make one constraint positive
            constraints = tuple(1 if j == 0 else -j for j in range(num_constraints))
        else:
            constraints = tuple(-j for j in range(num_constraints))
        attrs[i] = {_CONSTRAINTS_KEY: constraints}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(i) for i in range(num_trials)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 76.8μs -> 75.9μs (1.21% faster)
    for i, f in enumerate(feas):
        if i % 5 == 0:
            pass
        else:
            pass

def test_max_constraints_per_trial():
    # Trial with maximum allowed constraints (999)
    num_constraints = 999
    attrs = {0: {_CONSTRAINTS_KEY: tuple(-1 for _ in range(num_constraints))}}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(0)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 50.0μs -> 48.2μs (3.84% faster)

def test_max_trials():
    # 999 trials, each with 1 constraint
    num_trials = 999
    attrs = {i: {_CONSTRAINTS_KEY: (i % 2 - 1,)} for i in range(num_trials)}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(i) for i in range(num_trials)]
    vals, feas = _get_constraint_vals_and_feasibility(study, trials) # 327μs -> 321μs (1.93% faster)
    for i, f in enumerate(feas):
        pass

def test_large_scale_inconsistent_constraints_raises():
    # 500 trials, some with 10 constraints, some with 9
    num_trials = 500
    attrs = {}
    for i in range(num_trials):
        if i % 10 == 0:
            constraints = tuple(-j for j in range(9))
        else:
            constraints = tuple(-j for j in range(10))
        attrs[i] = {_CONSTRAINTS_KEY: constraints}
    study = DummyStudy(attrs)
    trials = [FrozenTrial(i) for i in range(num_trials)]
    with pytest.raises(ValueError, match="The number of constraints must be the same for all trials."):
        _get_constraint_vals_and_feasibility(study, trials) # 68.7μs -> 2.91μs (2261% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-_get_constraint_vals_and_feasibility-mhbenxu1 and push.

The optimization replaces the O(n²) list comprehension with validation approach with a single O(n) loop that builds the constraint list and validates lengths incrementally. **Key changes:** 1. **Single-pass processing**: Instead of creating the full `_constraint_vals` list first and then validating all lengths with `any()`, the optimized version processes trials one by one, checking constraint length consistency as it builds the list. 2. **Early validation**: The length check happens immediately when processing each trial, allowing for early exit on the first inconsistent constraint length rather than processing all trials before validation. 3. **Eliminated redundant iteration**: The original code iterates through trials twice - once in the list comprehension and once in the `any()` validation. The optimized version does both operations in a single loop. **Performance benefits:** - **Normal cases (4-11% faster)**: Eliminates the overhead of the `any()` generator expression and reduces total iterations from 2n to n - **Error cases with inconsistent constraints (60-2261% faster)**: Massive speedup because the function can exit immediately upon finding the first constraint length mismatch, rather than processing all trials first The optimization is particularly effective for large-scale scenarios where constraint length validation fails early, as shown in the test cases where inconsistent constraint detection becomes dramatically faster.

codeflash-ai bot requested a review from mashraf-222 October 29, 2025 02:57

codeflash-ai bot added ⚡️ codeflash Optimization PR opened by Codeflash AI 🎯 Quality: High Optimization Quality according to Codeflash labels Oct 29, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

⚡️ Speed up function `_get_constraint_vals_and_feasibility` by 9% #38

⚡️ Speed up function `_get_constraint_vals_and_feasibility` by 9% #38

Uh oh!

codeflash-ai bot commented Oct 29, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

⚡️ Speed up function _get_constraint_vals_and_feasibility by 9% #38

Are you sure you want to change the base?

⚡️ Speed up function _get_constraint_vals_and_feasibility by 9% #38

Uh oh!

Conversation

codeflash-ai bot commented Oct 29, 2025

📄 9% (0.09x) speedup for _get_constraint_vals_and_feasibility in optuna/samplers/_gp/sampler.py

📝 Explanation and details

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

⚡️ Speed up function `_get_constraint_vals_and_feasibility` by 9% #38

⚡️ Speed up function `_get_constraint_vals_and_feasibility` by 9% #38

📄 9% (0.09x) speedup for `_get_constraint_vals_and_feasibility` in `optuna/samplers/_gp/sampler.py`